from typing import Union, Dict

import numpy as np


class Rescaling:
    """
    Randomly rescale features of the sample.
    """
    def __init__(self, sigma: float = 0.5, keys=None) -> None:
        self.sigma = sigma
        self.keys = keys

    def _apply(self, sample: np.ndarray) -> np.ndarray:
        scaling_factor = np.random.normal(loc=1.0, scale=self.sigma, size=(1,)).astype(np.float32)
        return sample * scaling_factor
    
    def __call__(self, data: Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]) -> Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]:
        if isinstance(data, dict):
            for key in self.keys:
                sample = data[key]
                data[key] = self._apply(sample)
        elif isinstance(data, np.ndarray):
            data = self._apply(data)
        return data
